"""
Ollama 设置助手脚本 (demo/setup_ollama.py)

帮助用户检查和配置 Ollama 用于图像标签生成。
"""
import os
import sys
import logging
import subprocess
import requests
import platform

# 把项目根目录添加到 Python 路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
from config.config import APP_CONFIG

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def check_ollama_installed():
    """检查 Ollama 是否已安装"""
    try:
        result = subprocess.run(["ollama", "version"], capture_output=True, text=True)
        if result.returncode == 0:
            return True, result.stdout.strip()
        return False, None
    except FileNotFoundError:
        return False, None

def check_ollama_running(host="localhost", port="11434"):
    """检查 Ollama 服务是否正在运行"""
    try:
        response = requests.get(f"http://{host}:{port}/api/tags")
        if response.status_code == 200:
            return True, response.json()
        return False, None
    except requests.exceptions.ConnectionError:
        return False, None

def check_model_available(model="deepseek-vision", host="localhost", port="11434"):
    """检查模型是否已下载"""
    try:
        response = requests.get(f"http://{host}:{port}/api/tags")
        if response.status_code == 200:
            models = response.json()
            for model_info in models.get("models", []):
                if model_info.get("name") == model:
                    return True, model_info
            return False, None
        return False, None
    except requests.exceptions.ConnectionError:
        return False, None

def get_install_instructions():
    """获取安装 Ollama 的操作系统特定说明"""
    system = platform.system().lower()
    
    if system == "darwin":
        return "访问 https://ollama.com/download 下载 macOS 安装包"
    elif system == "linux":
        return "在终端运行: curl -fsSL https://ollama.com/install.sh | sh"
    elif system == "windows":
        return "访问 https://ollama.com/download 下载 Windows 安装包"
    else:
        return "访问 https://ollama.com/download 获取安装指南"

def check_and_setup():
    """检查 Ollama 状态并提供设置指南"""
    logger.info("检查 Ollama 配置...")
    
    # 检查是否安装了 Ollama
    installed, version = check_ollama_installed()
    if installed:
        logger.info(f"✅ Ollama 已安装: {version}")
    else:
        logger.error("❌ 未检测到 Ollama")
        logger.info(f"安装指南: {get_install_instructions()}")
        return
    
    # 获取配置
    load_dotenv()
    host = os.environ.get("OLLAMA_HOST", "http://localhost").replace("http://", "")
    port = os.environ.get("OLLAMA_PORT", "11434")
    model = os.environ.get("OLLAMA_MODEL", "deepseek-vision")
    
    # 检查 Ollama 服务是否运行
    running, models = check_ollama_running(host, port)
    if running:
        logger.info(f"✅ Ollama 服务正在运行 (http://{host}:{port})")
    else:
        logger.error(f"❌ Ollama 服务未运行 (http://{host}:{port})")
        logger.info("启动 Ollama 服务: 打开新终端窗口，运行命令 'ollama serve'")
        return
    
    # 检查是否已下载模型
    model_available, model_info = check_model_available(model, host, port)
    if model_available:
        logger.info(f"✅ 模型 {model} 已可用")
        if model_info:
            logger.info(f"   大小: {model_info.get('size', 'N/A')}")
    else:
        logger.warning(f"⚠️ 模型 {model} 未下载")
        logger.info(f"下载模型: 运行命令 'ollama pull {model}'")
    
    # 检查环境配置
    use_ollama = os.environ.get("USE_OLLAMA", "false").lower() == "true"
    if use_ollama:
        logger.info("✅ .env 中已配置使用 Ollama (USE_OLLAMA=true)")
    else:
        logger.warning("⚠️ .env 中未配置使用 Ollama (USE_OLLAMA=false)")
        logger.info("需要在 .env 文件中设置 USE_OLLAMA=true 以启用 Ollama")
    
    logger.info("\n总结:")
    logger.info("1. Ollama 安装状态: " + ("已安装 ✅" if installed else "未安装 ❌"))
    logger.info("2. Ollama 服务状态: " + ("运行中 ✅" if running else "未运行 ❌"))
    logger.info("3. 模型下载状态: " + ("已下载 ✅" if model_available else "未下载 ⚠️"))
    logger.info("4. 环境配置状态: " + ("已配置 ✅" if use_ollama else "未配置 ⚠️"))
    
    if installed and running and model_available and use_ollama:
        logger.info("\n🎉 一切配置正常，可以使用 Ollama 进行图像标签生成!")
        logger.info(f"测试命令: python3 demo/test_ollama_tagging.py --image [图像路径]")
    else:
        logger.info("\n⚠️ 需要完成上述步骤后才能使用 Ollama 进行图像标签生成")

if __name__ == "__main__":
    check_and_setup() 